import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread

from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class Form(HasTraits):
    frequencies = List(Float, [10])
    elevations = List(Float, [8])
    azimuthStart = Float(0)
    azimuthStop = Float(45)
    azimuthNumber = Int(46)
    po = Enum('HH', 'VV')
    environment = Enum('林地', '荒漠', '公路', '草地', '干雪', '潮雪', '湿雪', '砂土路')
    epsilon = Property(Complex, depends_on='environment')
    mu = Property(Complex, depends_on='environment')
    plot = Instance(Plot)
    progress = Str('0%')
    thread = QThread()

    def _get_epsilon(self):
        match self.environment:
            case '林地':
                return 3.9 - 1.9j
            case '荒漠':
                return 6.62 - 5.56j
            case '公路':
                return 4.9 - 0.23j
            case '草地':
                return 3.8 - 1.8j
            case '干雪':
                return 1.3 - 0.0002j
            case '潮雪':
                return 1.4 - 0.0002j
            case '湿雪':
                return 6 - 0.0002j
            case '砂土路':
                return 3.4 - 0.022j

    def _get_mu(self):
        return 1

    def calculate(self):
        v = self.v
        f = self.f
        frequencies = np.array(self.frequencies)
        elevations = np.array(self.elevations)
        azimuths = np.linspace(self.azimuthStart, self.azimuthStop, self.azimuthNumber)
        self.rcs = np.full((2, frequencies.size, elevations.size, azimuths.size), np.nan)  # 0 for hh, 1 for vv
        np.seterr(divide='ignore')
        v = np.concatenate((v[f[:, 0]][:, np.newaxis, :], v[f[:, 1]][:, np.newaxis, :], v[f[:, 2]][:, np.newaxis, :]), axis=1)
        e = np.roll(v, -1, axis=1) - v
        ec = (np.roll(v, -1, axis=1) + v) / 2
        n = np.cross(e[:, 0], e[:, 1])
        a = np.linalg.norm(n, axis=-1) / 2
        n = n / a[:, np.newaxis] / 2
        eps = 1e-8
        totalNumber = frequencies.size * elevations.size * azimuths.size
        for ff, frequency in enumerate(frequencies):
            k = 2 * np.pi / .3 * frequency
            for ee, elevation in enumerate(elevations):
                cosel = np.cos(elevation * np.pi / 180)
                sinel = np.sin(elevation * np.pi / 180)
                self.plot.title = f'Frequency={frequency}GHz,Elevation={elevation} degree'
                for aa, azimuth in enumerate(azimuths):
                    cosaz = np.cos(azimuth * np.pi / 180)
                    sinaz = np.sin(azimuth * np.pi / 180)
                    s = np.array([cosel * cosaz, cosel * sinaz, sinel])
                    # for HH polarization
                    i = np.array([-cosel * cosaz, -cosel * sinaz, -sinel])
                    w = s - i
                    ni = n @ i
                    light = ni < -eps
                    es = np.array([-sinaz, cosaz, 0])  ###HH
                    ei = np.array([-sinaz, cosaz, 0])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    nw = np.cross(n, w, axis=-1)
                    nw2 = (nw * nw).sum(-1)
                    nwSmall = nw2 < eps
                    perpendicular = nwSmall & light
                    oblique = (~nwSmall) & light
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * v[perpendicular, 0, :] @ w) * a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    # for scattering from ground
                    i = np.array([-cosel * cosaz, -cosel * sinaz, sinel])
                    w = s - i
                    ni = n @ i
                    light = ni < -eps
                    es = np.array([-sinaz, cosaz, 0])  ###HH
                    ei = np.array([-sinaz, cosaz, 0])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    nw = np.cross(n, w, axis=-1)
                    nw2 = (nw * nw).sum(-1)
                    nwSmall = nw2 < eps
                    perpendicular = nwSmall & light
                    oblique = (~nwSmall) & light
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * v[perpendicular, 0, :] @ w) * a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    z=self.mu*sinel/np.sqrt(self.epsilon*self.mu-cosel**2)
                    reflectionCoefficient=(z-1)/(z+1)
                    po += reflectionCoefficient*(poPerpendicular + poOblique)
                    self.rcs[0, ff, ee, aa] = -40 if np.isinf(tmp:=20 * np.log10(np.abs(po))) else tmp
                    # for VV polarization
                    i = np.array([-cosel * cosaz, -cosel * sinaz, -sinel])
                    w = s - i
                    ni = n @ i
                    light = ni < -eps
                    es = np.array([-sinel * cosaz, -sinel * sinaz, cosel])  ###VV
                    ei = np.array([-sinel * cosaz, -sinel * sinaz, cosel])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    nw = np.cross(n, w, axis=-1)
                    nw2 = (nw * nw).sum(-1)
                    nwSmall = nw2 < eps
                    perpendicular = nwSmall & light
                    oblique = (~nwSmall) & light
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * v[perpendicular, 0, :] @ w) * a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    # for scattering from ground
                    i = np.array([-cosel * cosaz, -cosel * sinaz, sinel])
                    w = s - i
                    ni = n @ i
                    light = ni < -eps
                    es = np.array([-sinel * cosaz, -sinel * sinaz, cosel])  ###VV
                    ei = np.array([-sinel * cosaz, -sinel * sinaz, cosel])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    nw = np.cross(n, w, axis=-1)
                    nw2 = (nw * nw).sum(-1)
                    nwSmall = nw2 < eps
                    perpendicular = nwSmall & light
                    oblique = (~nwSmall) & light
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * v[perpendicular, 0, :] @ w) * a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    z=self.epsilon*sinel/np.sqrt(self.epsilon*self.mu-cosel**2)
                    reflectionCoefficient=(z-1)/(z+1)
                    po += -reflectionCoefficient*(poPerpendicular + poOblique)
                    self.rcs[1, ff, ee, aa] = -40 if np.isinf(tmp:=20 * np.log10(np.abs(po))) else tmp
                    self.plot.data.update_data(
                        x=azimuths[:aa + 1],
                        hh=self.rcs[0, ff, ee, :aa + 1],
                        vv=self.rcs[1, ff, ee, :aa + 1],
                    )
                    self.progress = f'{((ff * elevations.size + ee) * azimuths.size + aa + 1) / totalNumber:<.1%}'
                    time.sleep(.1)
    def traits_init(self):
        self.thread.run = self.calculate
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            x=[],
            hh=[],
            vv=[],
        )
        self.plot.title = 'Ground'
        self.plot.x_axis.title = "Azimuth/deg"
        self.plot.y_axis.title = "RCS/dBsm"
        self.plot.plot(("x", "hh"), name='HH', color='auto')
        self.plot.plot(("x", "vv"), name='VV', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            UItem('progress'),
            spring
        ),
        width=1200,
        height=900,
        resizable=True,
    )

    azimuths_view = View(
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        resizable=True,
        buttons=OKCancelButtons,
    )

    frequencies_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )

    elevations_view = View(
        Group(
            UItem('elevations'),
            label='高低角（需要大于0°/°',
            show_border=True,
        ),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )

    environments_view = View(
        Item('environment', label='环境'),
        Item('epsilon', label='介电常数'),
        Item('mu', label='磁导率'),
        buttons=OKCancelButtons,
    )

